# 1. Tree models are hierarchical data structures that support parent-child relationships.

# 2. Let's add some icons to our screen graph viewer. We will use the Resource compiler that ships with PySide to compile our resources into
# binary and then import them into our application. First we need some icons. We have them in the project folder. Then we make a qrc
# file icons.qrc which contain xml to tell the compiler which icons we want to make into a binary. Once we have the xml file, we can
# go on and compile it. For that we need to send this qrc file to the PyRcc executable which compiles the files into binary. 
# Use this command for python 3: pyside-rcc -py3 icons.qrc -o icons_rc.py
# For python 2, remove the -py3 switch.
# Once we have the binary file, which contains the icons in binary format. All we need to do is to import that into our application,
# to use the icons in our scene graph viewer. The file should be placed in the same folder as the source code.

# 3. We have implemented our own Node class to store the data.
# 4. We also learned that, since tree models are hierarchical, they require us to implement two additional methods from the 
# QAbstractModelIndex class, the 'parent()' and the 'index()'.  

import sys
from PySide import QtCore, QtGui
import icons_rc                 # Import the icons file that was made earlier

class Node(object):
    def __init__(self, name, parent = None):
        self._name = name       # Name of the node
        self._children = []     # List of Children of the node
        self._parent = parent   # Parent of the node
        
        # If there is a parent to this node, then this node must add itself to the child list of the parent node
        if parent is not None:
            parent.addChild(self)
        
    # typeInfo method will return a string which will describe the type of this item. Since it is a node, it will return type as Node
    def typeInfo(self):
        return "Node"
    
    # Add a new child to the list of children
    def addChild(self, child):
        self._children.append(child)    # Append the new child to the list of children using append method.
        
    # Similar to the add child method above but puts the child at a given arbitrary location.
    def insertChild(self, position, child):
        # We check if the position is within the boundary of the list. If not, we return false.
        if position < 0 or position > len(self._children):
            return False
        # Else put the child at the given position using the insert method. Set the parent of the child to self and return true 
        self._children.insert(position, child)
        child._parent = self
        return True
    
    # Just like the insert child method, the remove child method is implemented as follows.
    def removeChild(self, position):
        # We check if the position is within the boundary of the list. If not, we return false.
        if position < 0 or position > len(self._children):
            return False
        # Else get the child at the given position using the pop method. Then et the parent of the child to None and return true 
        child = self._children.pop(position)
        child._parent = None
        return True    

        
    # Method to return the name of an object
    def name(self):
        return self._name
    
    # Method to set the node name as required by the setData method
    def setName(self, name):
        self._name = name
    
    # Method to get the child corresponding to an index (or row)
    def child(self, row):
        return self._children[row]
    
    # Method to get the number of children
    def childCount(self):
        return len(self._children)
    
    # Method to get the parent of this node
    def parent(self):
        return self._parent
    
    # Method to get the index number of this node relative to its parent
    def row(self):
        if self._parent is not None:
            return self._parent._children.index(self)
        
    # Method to display the data in the console
    def log(self, tabLevel = -1):
        output = ""
        tabLevel += 1
        
        for i in range(tabLevel):
            output += "\t"
        
        output += self._name + "\n"
        
        for child in self._children:
            output += child.log(tabLevel)
        
        tabLevel -= 1
        return output
    
    def __repr__(self):
        return self.log()
    
    
# Let's make some subclasses from the Node class and call them LightNode, TransformNode, CameraNode. All of them would need to 
# reimplement the typeInfo method. Initially we forced each class to have a parent. But now we changed it be without a parent.
class LightNode(Node):
    def __init__(self, name, parent = None):
        # Call the super class constructor
        super(LightNode, self).__init__(name, parent)
        
    # Reimplement the typeInfo method
    def typeInfo(self):
        return "LightNode"

class TransformNode(Node):
    def __init__(self, name, parent = None):
        # Call the super class constructor
        super(TransformNode, self).__init__(name, parent)
    # Reimplement the typeInfo method
    def typeInfo(self):
        return "TransformNode"

class CameraNode(Node):
    def __init__(self, name, parent = None):
        # Call the super class constructor
        super(CameraNode, self).__init__(name, parent)
    # Reimplement the typeInfo method
    def typeInfo(self):
        return "CameraNode"


# Now we will implement the tree model. The implement the tree model, we have to subclass QAbstractItemModel and provide methods for it.
class SceneGraphModel(QtCore.QAbstractItemModel):
    """ Inputs: Node, QObject """
    def __init__(self, root, parent = None):
        super(SceneGraphModel, self).__init__(parent)
        self._rootNode = root
    
    """ Inputs: QModelIndex """
    """ Outputs: int """
    """ Return the amount of children an item has """
    """ Note that we don't need the parent argument in the 'list' and the 'table' model. Only for 'tree' models. """
    """ We get the rootNode if the parent is not valid, otherwise we get the parent. Then we just return the childCount of the node """
    def rowCount(self, parent):
        if not parent.isValid():
            parentNode = self._rootNode
        else:
            parentNode = parent.internalPointer()
        
        return parentNode.childCount()  
    
    """ Inputs: QModelIndex """
    """ Outputs: int """
    """ Column count controls the number of columns in the tree view. Giving a column count of 2 makes two columns in the header of the tree view """
    """ But it does not display the other sections in tree view style, that's because the parent method of our model returns 0 on the second """
    """ parameter, which is actually the column of the createIndex method. This is telling that make column 0 (the first column) a hierarchical """ 
    """ view only. Now how should we use the second column. We can use the second column to display the type of the node. """ 
    """ Sounds like a good idea. For that we have to change the data and the headerData methods """
    def columnCount(self, parent):
        return 2
        
    # Data method to display something on the view. We will return both text and icons through this method.
    """ Inputs: QModelIndex, int """
    """ Outputs: QVariant, strings are cast to QString which is a QVariant """
    # Remember that the data method returns something that the view wants to display. All we will do here is call interalPointer to get our node,
    # then if it is displayrole, return the node name.
    # If the index is valid, we pass the test. Get the node. If it's displayrole, return the node name.
    # The index is a QModelIndex and it stores information about where an item is. Column corresponds to the header section.
    def data(self, index, role):
        if not index.isValid():
            return None
        
        node = index.internalPointer()
        
        # What to do when the view wants to display data, is done thorough DisplayRole.
        # Similarly, what to do in case when the view wants to edit data is done through the EditRole.
        if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole: # to retain value when item is double clicked for editing
            if index.column() == 0:     # if item lies in zeroth column then return the node name
                return node.name()
            else:                       # otherwise if item is not in zeroth column, return the node 'typeinfo'
                return node.typeInfo()  # to display the node typeInfo in the second column
        
        # Icons are returned by using a decoration role.
        if role == QtCore.Qt.DecorationRole:
            if index.column() == 0:     # if item lies in zeroth column then first get the 'typeinfo'. Depending on that, return the icon.
                typeInfo = node.typeInfo()
                if typeInfo == "LightNode":
                    # For returning an icon, we have to construct one using the QtGui.QIcon() method. The QIcon takes a QPixmap as an input.
                    # So we construct that as well. The QPixmap, takes a resource as input. To retrieve a resource at runtime, we type
                    # :/resource_name. Note that we do not want to work with full paths here.
                    return QtGui.QIcon(QtGui.QPixmap(":/Light.png"))
                if typeInfo == "CameraNode":
                    return QtGui.QIcon(QtGui.QPixmap(":/Camera.png"))
                if typeInfo == "TransformNode":
                    return QtGui.QIcon(QtGui.QPixmap(":/Transform.png"))
    
    # HeaderData method to display something on the view's header
    """ Inputs: int, Qt::Orientation, int """
    """ Outputs: QVariant, strings are cast to QString which is a QVariant """
    """ Let's just return the string "SceneGraph" to display on the header """
    def headerData(self, section, orientation, role):
        if role == QtCore.Qt.DisplayRole:
            if section == 0:
                return "SceneGraph"
            else:
                return "TypeInfo"   # To show TypeInfo in the header
    
    # Flags method to select the items, or enable items so that they are not greyed out etc. 
    """ Inputs: QModelIndex """
    """ Outputs: int (flag) """
    """ Flags method, as in previous examples, needs to return that the items are enabled, the items are selectable and the items are editable """
    """ The itemIsEditable flag requires that we implement a method called setData """
    def flags(self, index):
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable
    
    """ Inputs: QModelIndex, QVariant, int (flag) """
    """ The setData method receives an index (QModelIndex), a value (QVariant - A QVariant is a data holder for Qt applications) and a role """
    """ We want the role to default to EditRole """
    """ We will use the internalPointer to get our node and then set its name """
    def setData(self, index, value, role = QtCore.Qt.EditRole):
        if index.isValid():
            if role == QtCore.Qt.EditRole:
                node = index.internalPointer()
                node.setName(value)
                return True
        return False
    
    
    # Tree models and Tree view require two more special methods
    """ Remember that the QModelIndex class has valuable information about where and item was with the row and columns methods that return ints"""
    """ In the previous examples, we used those methods index our lists. """
    """ The QModelIndex also have something called and index pointer. The indexPointer returns a custom node instance which we implemented above """
    # Our goal of from this method is to return the parent of the given index and the returned value must be a QModelIndex class itself
    """ Inputs: QModelIndex """
    """ Outputs: QModelIndex """
    """ Returns the parent of the node with the given QModelIndex """
    def parent(self, index):
        node = self.getNode(index) # Returns a node
        parentNode = node.parent()
        # If the returned node is the 'root' node then return an empty QModelIndex because roots don't have a parent
        if parentNode == self._rootNode:
            return QtCore.QModelIndex()
        # If the returned node is not the root node, then it must be a child node. Now is the time to create a QModelIndex class object that
        # we will wrap the parent node into. This is done by the factor method called 'createIndex'. We have this method because of inheriting
        # from the QAbstractItemModel. First arg is an integer that corresponds to the row that the parent node lies in, so we implement the
        # custom method row() that we created above. The second arg can be 0 for now. The third argument is an object that can be anything. It's
        # also the object that will be returned when you call internal pointer of the index that this method returns. We will pass our parent node
        # to this argument.
        # self.createIndex(int, int, object=0)
        return self.createIndex(parentNode.row(), 0, parentNode)

    """ Inputs: int, int, QModelIndex """
    """ Outputs: QModelIndex """
    """ Returns the QModelIndex that corresponds to the given row, column and parent node """
    # The row and column are integers that the view acquires from the rowCount and the columnCount methods. It also receives a parent which it
    # acquires from the method implemented above. The index method is responsible for returning a child at the given row and given column of the
    # given parent.   
    def index(self, row, column, parent):
        parentNode = self.getNode(parent)
        
        # Now we have the parent node. After that, we ask the parentNode to get the child by calling our own child method and passing the row 
        # that we want the child from. The row is an integer.
        childItem = parentNode.child(row)
        
        # But if for some wierd reason the child is None, we return an empty QModelIndex. If child exists we create an index and return it.
        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QtCore.QModelIndex

    # Let's make a custom method, that is actually not overriding anything.
    # It will act like a helper. If we pass a QModelIndex with a valid index, we will get the internal pointer (i.e. the node,
    # if it exists) and return it. Otherwise we return the root node.  
    # First check if the parent is valid. Because in the 'parent' method above notice that we return an empty QModelIndex in case of a root
    # node, which is not valid. Therefore we set the parent node here to rootNode if the parent is not valid i.e. an empty QModelIndex. 
    # But if parent is valid then we get the internal pointer which again corresponds to our node instance.
    """ Custom method (not overriding anything) """
    """ Inputs: QModelIndex """
    def getNode(self, index):
        if index.isValid():
            node = index.internalPointer()
            if node:
                return node
        else:
            return self._rootNode
        
    # Methods for Inserting and Removing rows at runtime.
    # Recall that in the last examples, we ignored the parent arguments. Now is the time we will use it.
    # Insert rows acts as an insert children method, i.e. you are basically inserting children to this parent.
    # For that we will make an insert child method in the Node class which will add a child at an arbitrary location. 
    """ Inputs: int, int, QModelIndex """
    def insertRows(self, position, rows, parent = QtCore.QModelIndex()):
        # To insert the children, we need to get the parent node and we get it by calling our customly implemented getNode method
        # by passing the parent QModelIndex class to it. As a result it returns its internalPointer which is the name of the node.
        parentNode = self.getNode(parent)
    # Inputs are parent to insert at, the position to insert at, and how many items (more accurately, the ending index of items)
        self.beginInsertRows(parent, position, position + rows - 1)
        for row in range(rows):
            # Keep track of child count
            childCount = parentNode.childCount()
            # Create a node
            childNode = Node("NewNode" + str(childCount))
            # Insert this child node. The current child node pushes the previous child nodes, downwards 
            success = parentNode.insertChild(position, childNode)
        self.endInsertRows()
        return success
    
    """ Inputs: int, int, QModelIndex """
    def removeRows(self, position, rows, parent = QtCore.QModelIndex()):
        parentNode = self.getNode(parent)
        self.beginRemoveRows(parent, position, position + rows - 1)
        for row in range(rows):
            success = parentNode.removeChild(position)
        self.endRemoveRows()
        return success
    
    # As a demonstration we can also insert light nodes
    def insertLightRows(self, position, rows, parent = QtCore.QModelIndex()):
        parentNode = self.getNode(parent)
        self.beginInsertRows(parent, position, position + rows - 1)
        for row in range(rows):
            # Keep track of child count
            childCount = parentNode.childCount()
            # Create a node
            childNode = LightNode("NewLightNode" + str(childCount))
            # Insert this child node. The current child node pushes the previous child nodes, downwards 
            success = parentNode.insertChild(position, childNode)
        self.endInsertRows()
        return success
    
        
if __name__ == "__main__":
    
    appObject = QtGui.QApplication(sys.argv)
    appObject.setStyle('plastique')
    
    rootNode = Node("Head")
    childNode1 = Node("Torso", rootNode)
    childNode2 = Node("RightArm", childNode1)
    childNode3 = CameraNode("LeftArm", childNode1)
    childNode4 = Node("RightWrist", childNode2)
    childNode5 = TransformNode("LeftWrist", childNode3)
    childNode6 = Node("RightHand", childNode2)
    childNode7 = Node("LeftHand", childNode3)
    childNode8 = LightNode("RightElbow", childNode2)
    childNode8 = Node("LeftElbow", childNode3)
    
    print(rootNode)
    
    # Create our Scenegraph model which we created by providing it the rootNode, which it stores in its protected member _rootNode
    model = SceneGraphModel(rootNode)
    
    # Create a treeview and show it
    treeview = QtGui.QTreeView()
    treeview.show()
    
    # Finally set the treeview model using setModel()
    treeview.setModel(model)
    
    # Insert some items at the root
    model.insertRows(0, 5)
    
    # Insert some light nodes at index location 1
    model.insertLightRows(1, 5)
    
    # Remove some items from the top
    model.removeRows(0, 2)
    
    # Insert some rows under the LeftArm child (i.e. under a node that already exists).
    # insertRows and removeRows methods have a parent parameter that takes a QModelIndex argument. If you don't specify it,
    # it defaults to root. But we will use the parent argument this time.
    # To specify that, we need a QModelIndex that corresponds to a child. Let's get the LeftArm node (childNode3).
    # For that we use the index() method. Get the item at row=0, column=0 and we specify and an empty QModelIndex (i.e. the root).
    # First we get to the torso. Since the root is Head, it will return the Torso to us.  
    Torso = model.index(8, 0, QtCore.QModelIndex())
    # Then to get the left arm, leftArm is the second row (index 1) under the torso, i.e. rightArm=0, leftArm=1
    LeftArm = model.index(1, 0, Torso)
    # Insert rows under the left arm
    model.insertRows(0, 5, LeftArm)
    
    # Execute the application
    sys.exit(appObject.exec_())


 